// by mircemk, July 2026

// ============================================================
//  Audio Spectrum Analyzer V2.1 TUNED  (Single FFT)
//  Hardware : Elecrow CrowPanel ESP32 3.5" (480x320)
//  ADC Input: GPIO35, DC bias ~1.6V (2x100k divider)
//
//  TUNED VERSION:
//  - Single FFT over whole spectrum
//  - 1024 samples @ 40kHz  -> much faster than 2048
//  - 24 bars -> better practical resolution on ESP32
//  - Tuned band energy formula
//  - Mild HF tilt compensation
//  - Mild mid cut around 500 Hz – 2 kHz
//  - Noise floor calibration once at startup
// ============================================================

#include <Arduino.h>
#include <TFT_eSPI.h>
#include <arduinoFFT.h>
#include <math.h>

// ─────────────────────────────────────────────
//  Екран
// ─────────────────────────────────────────────
#define SCREEN_WIDTH    480
#define SCREEN_HEIGHT   320
#define MARGIN           35
#define NUM_BARS         24
#define BAR_GAP           2
#define BACKGROUND      TFT_BLACK
#define SCALE_COLOR     0x07FF

// ─────────────────────────────────────────────
//  ADC
// ─────────────────────────────────────────────
#define ADC_PIN          35

// ─────────────────────────────────────────────
//  FFT
//  1024 @ 40kHz -> 39.06 Hz/bin
// ─────────────────────────────────────────────
#define SAMPLES_FFT      1024
#define SAMPLING_FREQ   40000

// ─────────────────────────────────────────────
//  Spectrum range
// ─────────────────────────────────────────────
#define FREQ_MIN        20.0f
#define FREQ_MAX     20000.0f

// ─────────────────────────────────────────────
//  Noise floor
// ─────────────────────────────────────────────
#define NOISE_CAL_FRAMES  12
#define NOISE_MARGIN       1.5f

// ─────────────────────────────────────────────
//  Visual tuning
// ─────────────────────────────────────────────
#define PEAK_HOLD_FRAMES  10
#define FALL_SPEED         25
#define PEAK_FALL_SPEED    12
#define RISE_SMOOTH      0.35f

// ─────────────────────────────────────────────
//  Globals
// ─────────────────────────────────────────────
TFT_eSPI tft = TFT_eSPI();

double vReal[SAMPLES_FFT];
double vImag[SAMPLES_FFT];

ArduinoFFT<double> FFT = ArduinoFFT<double>(vReal, vImag, SAMPLES_FFT, SAMPLING_FREQ);

uint16_t barPalette[NUM_BARS];
float    noiseFloor[NUM_BARS];

float barHeight     [NUM_BARS];
int   prevBarHeight [NUM_BARS];
int   peakHeight    [NUM_BARS];
int   prevPeakHeight[NUM_BARS];
int   peakHoldCount [NUM_BARS];

int displayAreaWidth;
int displayAreaHeight;
int startX;
float barWidth;
int maxBarHeight;

int bandLo[NUM_BARS];
int bandHi[NUM_BARS];

// ─────────────────────────────────────────────
//  slotColor()
// ─────────────────────────────────────────────
uint16_t slotColor(int pixelFromBottom) {
    uint8_t g = (uint8_t)map(pixelFromBottom, 0, displayAreaHeight, 0, 30);
    return tft.color565(g, g, g);
}

// ─────────────────────────────────────────────
//  createPalette()
// ─────────────────────────────────────────────
void createPalette() {
    struct RGB { uint8_t r, g, b; };
    RGB colors[] = {
        {255,   0,   0},
        {255, 255,   0},
        {  0, 255,   0},
        {  0, 255, 255},
        {  0,   0, 255},
        {255,   0, 255}
    };

    for (int i = 0; i < NUM_BARS; i++) {
        float pos = (float)i / (NUM_BARS - 1) * 5.0f;
        int idx = (int)pos;
        float frac = pos - idx;

        uint8_t r, g, b;
        if (idx >= 5) {
            r = colors[5].r;
            g = colors[5].g;
            b = colors[5].b;
        } else {
            r = colors[idx].r + (colors[idx + 1].r - colors[idx].r) * frac;
            g = colors[idx].g + (colors[idx + 1].g - colors[idx].g) * frac;
            b = colors[idx].b + (colors[idx + 1].b - colors[idx].b) * frac;
        }

        barPalette[i] = tft.color565(r, g, b);
    }
}

// ─────────────────────────────────────────────
//  showIntroText()
// ─────────────────────────────────────────────
void showIntroText() {
    tft.fillScreen(TFT_BLACK);
    tft.setTextDatum(MC_DATUM);
    tft.setTextColor(TFT_WHITE, TFT_BLACK);
    tft.drawCentreString("Spectrum Analyzer", SCREEN_WIDTH / 2, 95, 4);
    tft.drawCentreString("by",               SCREEN_WIDTH / 2, 145, 4);
    tft.drawCentreString("mircemk",          SCREEN_WIDTH / 2, 190, 4);
    delay(3000);
}

// ─────────────────────────────────────────────
//  drawTubeBar()
// ─────────────────────────────────────────────
void drawTubeBar(int x, int y, int w, int h, uint16_t baseColor) {
    if (h <= 0 || w <= 0) return;

    uint8_t r_b = (baseColor >> 11) & 0x1F;
    uint8_t g_b = (baseColor >>  5) & 0x3F;
    uint8_t b_b =  baseColor        & 0x1F;

    for (int i = 0; i < w; i++) {
        float intensity = 0.5f + 0.5f * sinf((PI * i) / (w - 1));
        uint16_t litColor = tft.color565(
            (uint8_t)(r_b * intensity) << 3,
            (uint8_t)(g_b * intensity) << 2,
            (uint8_t)(b_b * intensity) << 3
        );
        tft.drawFastVLine(x + i, y, h, litColor);
    }
}

// ─────────────────────────────────────────────
//  drawUI()
// ─────────────────────────────────────────────
void drawUI() {
    tft.fillScreen(BACKGROUND);

    tft.setTextDatum(TC_DATUM);
    tft.setTextColor(TFT_YELLOW, BACKGROUND);
    tft.drawString("Spectrum Analyzer [L]", SCREEN_WIDTH / 2, 5, 4);

    tft.setTextColor(SCALE_COLOR, BACKGROUND);
    const char* dbValues[] = {"+3","-3","-9","-15","-21","-27","-33","-39","-45"};
    float yStep = (float)(SCREEN_HEIGHT - 2 * MARGIN) / 8.0f;

    for (int i = 0; i < 9; i++) {
        int y = MARGIN + (int)(i * yStep);
        tft.drawString(dbValues[i], 18, y - 7, 2);
        tft.drawFastHLine(35, y, 8, SCALE_COLOR);
        tft.drawString(dbValues[i], SCREEN_WIDTH - 18, y - 7, 2);
        tft.drawFastHLine(SCREEN_WIDTH - 43, y, 8, SCALE_COLOR);
    }

    const char* freqValues[] = {"20","50","100","200","500","1K","2K","5K","10K","20K"};
    for (int i = 0; i < 10; i++) {
        int x = MARGIN + 20 + (int)(i * ((SCREEN_WIDTH - 2 * MARGIN - 40) / 9.0f));
        tft.drawString(freqValues[i], x, SCREEN_HEIGHT - MARGIN + 8, 2);
    }

    for (int i = 0; i < NUM_BARS; i++) {
        int xPos = startX + (int)(i * (barWidth + BAR_GAP));
        for (int j = 0; j < displayAreaHeight; j++) {
            tft.drawFastHLine(xPos, (SCREEN_HEIGHT - MARGIN) - j, (int)barWidth, slotColor(j));
        }
    }
}

// ─────────────────────────────────────────────
//  sampleADC()
// ─────────────────────────────────────────────
void sampleADC(double* buf, int nSamples, int freqHz) {
    unsigned long periodUs = 1000000UL / (unsigned long)freqHz;
    unsigned long t0 = micros();
    long dcSum = 0;

    for (int i = 0; i < nSamples; i++) {
        int raw = analogRead(ADC_PIN);
        buf[i] = (double)raw;
        dcSum += raw;
        while ((micros() - t0) < (unsigned long)(i + 1) * periodUs) { }
    }

    int dcMid = (int)(dcSum / nSamples);
    for (int i = 0; i < nSamples; i++) {
        buf[i] -= dcMid;
    }
}

// ─────────────────────────────────────────────
//  prepareBandBounds()
//  Log spaced bands
// ─────────────────────────────────────────────
void prepareBandBounds() {
    float res = (float)SAMPLING_FREQ / (float)SAMPLES_FFT;

    for (int b = 0; b < NUM_BARS; b++) {
        float fLo = FREQ_MIN * powf(FREQ_MAX / FREQ_MIN, (float)b / NUM_BARS);
        float fHi = FREQ_MIN * powf(FREQ_MAX / FREQ_MIN, (float)(b + 1) / NUM_BARS);

        int lo = max(1, (int)(fLo / res));
        int hi = min((SAMPLES_FFT / 2) - 1, (int)(fHi / res));

        if (hi < lo) hi = lo;

        // Avoid too many identical first bands
        if (b > 0 && lo <= bandHi[b - 1]) {
            lo = bandHi[b - 1] + 1;
            if (lo > hi) hi = lo;
            if (hi > (SAMPLES_FFT / 2) - 1) hi = (SAMPLES_FFT / 2) - 1;
            if (lo > hi) lo = hi;
        }

        bandLo[b] = lo;
        bandHi[b] = hi;
    }

    // Safety fix: ensure valid order
    for (int b = 0; b < NUM_BARS; b++) {
        if (bandLo[b] < 1) bandLo[b] = 1;
        if (bandHi[b] < bandLo[b]) bandHi[b] = bandLo[b];
        if (bandHi[b] > (SAMPLES_FFT / 2) - 1) bandHi[b] = (SAMPLES_FFT / 2) - 1;
    }
}

// ─────────────────────────────────────────────
//  runFFT()
//  Tuned aggregation:
//  - sumPower from bins
//  - divide by count^0.65  (compromise between sum and avg)
//  - mild HF tilt
//  - mild mid cut 500Hz..2kHz
// ─────────────────────────────────────────────
void runFFT(float* outRaw) {
    for (int i = 0; i < SAMPLES_FFT; i++) vImag[i] = 0.0;

    FFT.windowing(FFTWindow::Hamming, FFTDirection::Forward);
    FFT.compute(FFTDirection::Forward);
    FFT.complexToMagnitude();

    float binHz = (float)SAMPLING_FREQ / (float)SAMPLES_FFT;

    for (int b = 0; b < NUM_BARS; b++) {
        double sumPower = 0.0;
        int count = 0;

        for (int bin = bandLo[b]; bin <= bandHi[b]; bin++) {
            double mag = vReal[bin];
            sumPower += (mag * mag);
            count++;
        }

        float val = (count > 0)
            ? (float)(sqrt(sumPower) / pow((double)count, 0.65))
            : 0.0f;

        float pos = (float)b / (NUM_BARS - 1);     // 0..1
        float tilt = 1.0f + pos * 2.25f;           // 1.0 .. 2.25
        val *= tilt;

        float fCenter = sqrtf((bandLo[b] * binHz) * (bandHi[b] * binHz));
        if (fCenter > 500.0f && fCenter < 2000.0f) {
            val *= 0.42f;
        }

        outRaw[b] = val;
    }
}

// ─────────────────────────────────────────────
//  calibrateNoise()
// ─────────────────────────────────────────────
void calibrateNoise() {
    tft.setTextDatum(MC_DATUM);
    tft.setTextColor(TFT_CYAN, BACKGROUND);
    tft.drawCentreString("Calibrating...", SCREEN_WIDTH / 2, SCREEN_HEIGHT / 2 - 20, 4);
    tft.setTextColor(TFT_DARKGREY, BACKGROUND);
    tft.drawCentreString("Keep audio input silent", SCREEN_WIDTH / 2, SCREEN_HEIGHT / 2 + 15, 2);

    for (int i = 0; i < NUM_BARS; i++) noiseFloor[i] = 0.0f;

    float raw[NUM_BARS];

    for (int frame = 0; frame < NOISE_CAL_FRAMES; frame++) {
        sampleADC(vReal, SAMPLES_FFT, SAMPLING_FREQ);
        runFFT(raw);

        for (int i = 0; i < NUM_BARS; i++) {
            if (raw[i] > noiseFloor[i]) noiseFloor[i] = raw[i];
        }

        int pw = (SCREEN_WIDTH - 80) * (frame + 1) / NOISE_CAL_FRAMES;
        tft.fillRect(40, SCREEN_HEIGHT / 2 + 40, pw, 10, TFT_CYAN);
    }

    for (int i = 0; i < NUM_BARS; i++) {
        if (noiseFloor[i] < 8.0f) noiseFloor[i] = 8.0f;
    }

    drawUI();
}

// ─────────────────────────────────────────────
//  mapBandToHeight()
// ─────────────────────────────────────────────
int mapBandToHeight(float magnitude) {
    if (magnitude < 1.0f) return 0;

    float db = 20.0f * log10f(magnitude);

    // Slightly more aggressive than 75 dB
    float normalized = constrain(db / 95.0f, 0.0f, 1.0f);
    return (int)(normalized * maxBarHeight);
}

// ─────────────────────────────────────────────
//  updateBars()
// ─────────────────────────────────────────────
void updateBars(float* bands) {
    for (int i = 0; i < NUM_BARS; i++) {
        float targetH = (float)mapBandToHeight(bands[i]);

        if (targetH >= barHeight[i]) {
            barHeight[i] += (targetH - barHeight[i]) * RISE_SMOOTH;
            if ((targetH - barHeight[i]) < 0.5f) barHeight[i] = targetH;
        } else {
            barHeight[i] -= FALL_SPEED;
            if (barHeight[i] < 0.0f) barHeight[i] = 0.0f;
        }

        int barH_int = (int)barHeight[i];

        if (barH_int >= peakHeight[i]) {
            peakHeight[i] = barH_int;
            peakHoldCount[i] = PEAK_HOLD_FRAMES;
        } else {
            if (peakHoldCount[i] > 0) {
                peakHoldCount[i]--;
            } else {
                peakHeight[i] -= PEAK_FALL_SPEED;
                if (peakHeight[i] < 0) peakHeight[i] = 0;
            }
        }
    }
}

// ─────────────────────────────────────────────
//  redrawBars()
// ─────────────────────────────────────────────
void redrawBars() {
    int bottomY = SCREEN_HEIGHT - MARGIN;

    for (int i = 0; i < NUM_BARS; i++) {
        int xPos = startX + (int)(i * (barWidth + BAR_GAP));
        int w    = (int)barWidth;
        int newH = (int)barHeight[i];
        int oldH = prevBarHeight[i];
        int newP = peakHeight[i];
        int oldP = prevPeakHeight[i];

        if (newH > oldH) {
            drawTubeBar(xPos, bottomY - newH, w, newH - oldH, barPalette[i]);
        } else if (newH < oldH) {
            for (int px = newH; px < oldH; px++) {
                tft.drawFastHLine(xPos, bottomY - px - 1, w, slotColor(px));
            }
        }

        if (oldP > 0 && (oldP != newP || oldP <= newH)) {
            int oldPY = bottomY - oldP - 1;
            if (oldPY >= MARGIN && oldP > newH) {
                tft.drawFastHLine(xPos, oldPY, w, slotColor(oldP));
            }
        }

        if (newP > newH + 1) {
            int newPY = bottomY - newP - 1;
            if (newPY >= MARGIN) {
                tft.drawFastHLine(xPos, newPY, w, TFT_WHITE);
            }
        }

        prevBarHeight[i]  = newH;
        prevPeakHeight[i] = (newP > newH + 1) ? newP : 0;
    }
}

// ─────────────────────────────────────────────
//  setup()
// ─────────────────────────────────────────────
void setup() {
    Serial.begin(115200);

    analogSetAttenuation(ADC_11db);
    analogReadResolution(12);
    pinMode(ADC_PIN, INPUT);

    tft.begin();
    tft.setRotation(1);

    displayAreaWidth  = SCREEN_WIDTH  - (2 * MARGIN) - 30;
    displayAreaHeight = SCREEN_HEIGHT - (2 * MARGIN);
    startX            = MARGIN + 15;
    barWidth          = (float)(displayAreaWidth - (NUM_BARS - 1) * BAR_GAP) / NUM_BARS;
    maxBarHeight      = displayAreaHeight;

    createPalette();
    prepareBandBounds();
    showIntroText();
    drawUI();
    calibrateNoise();

    memset(barHeight,      0, sizeof(barHeight));
    memset(prevBarHeight,  0, sizeof(prevBarHeight));
    memset(peakHeight,     0, sizeof(peakHeight));
    memset(prevPeakHeight, 0, sizeof(prevPeakHeight));
    memset(peakHoldCount,  0, sizeof(peakHoldCount));
}

// ─────────────────────────────────────────────
//  loop()
// ─────────────────────────────────────────────
void loop() {
    float raw[NUM_BARS];
    float bands[NUM_BARS];

    sampleADC(vReal, SAMPLES_FFT, SAMPLING_FREQ);
    runFFT(raw);

    for (int i = 0; i < NUM_BARS; i++) {
        float sig = raw[i] - noiseFloor[i] * NOISE_MARGIN;
        bands[i] = (sig < 0.0f) ? 0.0f : sig;
    }

    updateBars(bands);
    redrawBars();
}